# =============================================================================
# Statistical Analysis Code for:
# "Chronic adaptive versus conventional DBS response patterns in Parkinson's
#  disease: A pilot randomized crossover trial"
#
# Exploratory Analysis: Treatment Effect Modifiers
# Author: Jun Tanimura, MD, MSc.
# Created: 2025-09-07
# Last Updated: 2025-09-07
#
# PURPOSE:
# This code identifies baseline patient characteristics that modify the treatment
# effect of aDBS versus cDBS. It examines whether patients with more higher baseline
# burden benefit more (or less) from aDBS compared to cDBS.
#
# KEY FEATURES:
# - Uses unified sign convention: positive effects always mean "worse baseline -> 
#   better outcome with aDBS" for clinical interpretability
# - Analyzes interactions between treatment (aDBS vs cDBS) and baseline modifiers
# - Accounts for crossover design (Period, Sequence effects) and baseline adjustment
# - Uses mixed-effects models to account for repeated measurements within patients
# - Calculates effect sizes (Cohen's d and partial r) for clinical relevance
# - We report only the result of partial correlation (partial r)
#
# IMPORTANT NOTES:
# - Time variables (ON time, OFF time, dyskinesia) are log-transformed for 
#   statistical modeling, but results are interpreted on original scale by back-transforming.
# - All baseline predictors are standardized (z-scores) for comparability
# - Multiple comparisons are performed (many modifier-outcome combinations);
#   results should be interpreted as exploratory/hypothesis-generating!
# - Robust effects are defined based on the LOO sensitivity analysis, which is
#   separately conducted.
# =============================================================================

library(tidyverse)
library(lme4)
library(lmerTest)
library(emmeans)

# =============================================================================
# Configuration: Define outcomes and their clinical interpretation
# =============================================================================

# Clinical outcomes to analyze
eval_items_ordered <- c("ADL", "UPDRS_part_1", "UPDRS_part_2", "UPDRS_part_3", "UPDRS_part_4",
                        "UPDRS_total", "ON_time", "OFF_time",
                        "Dys_sev_time", "Dys_weak_time", "MMSE", "PDQ_39")

# Outcomes where lower values indicate better clinical status
# For these, treatment effects will be reversed in sign for unified interpretation
lower_the_better_metrics <- c("UPDRS_part_1", "UPDRS_part_2", "UPDRS_part_3", "UPDRS_part_4", 
                              "UPDRS_total", "PDQ_39", "Dys_sev_time", 
                              "Dys_weak_time", "OFF_time", 
                              "age", "PD_duration", "P_HY", "LEDD")

# Outcomes where higher values indicate better clinical status
higher_the_better_metrics <- c("MMSE", "ON_time")

baseline_predictors <- c("age", "PD_duration", "sex", "LEDD",
                         "P_ADL", "P_MMSE", "P_HY",
                         "P_UPDRS_part_1", "P_UPDRS_part_2", "P_UPDRS_part_3", 
                         "P_UPDRS_part_4", "P_UPDRS_total",
                         "P_ON_time", "P_OFF_time", 
                         "P_Dys_sev_time", "P_Dys_weak_time",
                         "P_PDQ_39")

predictor_better_when_higher <- c("P_ADL", "P_MMSE", "P_ON_time")

predictor_worse_when_higher <- c("age", "PD_duration", "LEDD",
                                 "P_UPDRS_part_1", "P_UPDRS_part_2", "P_UPDRS_part_3", 
                                 "P_UPDRS_part_4", "P_UPDRS_total", "P_HY", "P_PDQ_39", 
                                 "P_OFF_time", "P_Dys_sev_time", 
                                 "P_Dys_weak_time")

# =============================================================================
# Data Processing and Preparation
# =============================================================================

# Read data and standardize column names (replace hyphens with underscores)
df <- read_csv("data.csv")
names(df) <- gsub("-", "_", names(df))

# Identify which baseline predictors are available in the dataset
existing_predictors <- baseline_predictors[baseline_predictors %in% names(df)]
df_sign_unified <- df

# =============================================================================
# Unified Sign Convention: Making results clinically interpretable
# =============================================================================
# The unified sign convention ensures that positive interaction effects always
# mean "patients with higher baseline burden benefit more from aDBS"
# This makes it easier to interpret results across different outcomes and modifiers
# regardless of whether higher or lower values are clinically better.

predictor_vars <- existing_predictors[!existing_predictors %in% c("sex")]

# Create unified versions: reverse sign for predictors where higher = better baseline
# Example: Higher baseline ADL = better baseline, so we reverse it to make
# "worse baseline" consistently represented by higher values in unified scale
for(pred in predictor_vars) {
  if(pred %in% names(df_sign_unified)) {
    original_values <- df_sign_unified[[pred]]
    unified_values <- if(pred %in% predictor_better_when_higher) -original_values else original_values
    df_sign_unified[[paste0(pred, "_unified")]] <- unified_values
  }
}

# Z-score standardization: Convert all predictors to same scale (mean=0, SD=1)
# This allows comparing effect sizes across different baseline measures
# (e.g., age in years vs. UPDRS scores) on the same scale
for(pred in predictor_vars) {
  unified_col <- paste0(pred, "_unified")
  if(unified_col %in% names(df_sign_unified)) {
    df_sign_unified[[paste0(pred, "_standardized")]] <- scale(df_sign_unified[[unified_col]])[,1]
  }
}

available_standardized_predictors <- paste0(predictor_vars, "_standardized")
if("sex" %in% existing_predictors) {
  available_standardized_predictors <- c(available_standardized_predictors, "sex")
}

# =============================================================================
# Transform Data to Long Format for Mixed-Effects Analysis
# =============================================================================
# Convert from wide format (one row per patient) to long format (one row per
# patient-treatment combination). This allows us to model treatment effects
# while accounting for within-patient correlation.

ancova_data <- map_dfr(eval_items_ordered, function(item) {
  col_P <- paste0("P_", item)
  col_C <- paste0("C_", item)
  col_A <- paste0("A_", item)
  
  if(item == "PDQ_39") {
    col_P <- "P_PDQ_39"
    col_C <- "C_PDQ_39"
    col_A <- "A_PDQ_39"
  }
  
  if(all(c(col_P, col_C, col_A) %in% names(df_sign_unified))) {
    df_sign_unified %>%
      select(PatientID, Group, all_of(c(col_P, col_C, col_A)), 
             all_of(available_standardized_predictors)) %>%
      pivot_longer(
        cols = c(all_of(col_C), all_of(col_A)),
        names_to = "Treatment_raw",
        values_to = "Score_post"
      ) %>%
      mutate(
        Baseline = !!sym(col_P),
        Treatment = case_when(
          str_detect(Treatment_raw, "^C_") ~ "cDBS",
          str_detect(Treatment_raw, "^A_") ~ "aDBS"
        ),
        Period = case_when(
          Group == 1 & Treatment == "cDBS" ~ "Phase1",
          Group == 1 & Treatment == "aDBS" ~ "Phase2", 
          Group == 2 & Treatment == "aDBS" ~ "Phase1",
          Group == 2 & Treatment == "cDBS" ~ "Phase2"
        ),
        Evaluation = item,
        # Identify time variables (ON time, OFF time, dyskinesia time)
        # These need log transformation because time values sometimes include 0.
        is_time_var = str_detect(item, "time"),
        # Log transformation: log1p(x) = log(1+x) handles zeros gracefully
        # For non-time variables, use original values
        Score_transformed = ifelse(is_time_var, log1p(pmax(Score_post, 0)), Score_post),
        Baseline_transformed = ifelse(is_time_var, log1p(pmax(Baseline, 0)), Baseline)
      ) %>%
      select(PatientID, Group, Evaluation, Treatment, Period,
             Baseline, Score_post, Baseline_transformed, Score_transformed,
             all_of(available_standardized_predictors))
  }
}) %>%
  filter(!is.na(Score_post)) %>%
  mutate(
    Treatment = factor(Treatment, levels = c("cDBS", "aDBS")),
    Period = factor(Period, levels = c("Phase1", "Phase2")),
    Evaluation = factor(Evaluation, levels = eval_items_ordered),
    PatientID = factor(PatientID)
  )

# Standardize baseline values within each outcome (z-score normalization)
# This ensures baseline adjustment accounts for different scales across outcomes
# (e.g., UPDRS scores vs. time in hours)
ancova_data <- ancova_data %>%
  group_by(Evaluation) %>%
  mutate(Baseline_z = as.numeric(scale(Baseline_transformed))) %>%
  ungroup()

if("sex" %in% existing_predictors) {
  ancova_data <- ancova_data %>% mutate(sex = factor(sex))
}

# =============================================================================
# Modifier Analysis Function
# =============================================================================
# This function tests whether a baseline characteristic modifies the treatment effect.
# It compares two models:
# 1. BASE MODEL: Treatment + Period + Sequence + Baseline + random patient effects
# 2. MODIFIER MODEL: Same as base + interaction between Treatment and modifier
#
# If the interaction is significant, it means the treatment effect differs
# depending on the baseline value of the modifier.

run_modifier_analysis <- function(data, metric, modifier_var) {
  # Filter data for this specific outcome and remove missing values
  metric_data <- data %>% 
    filter(Evaluation == metric) %>%
    filter(!is.na(Score_transformed) & !is.na(Baseline_z) & !is.na(!!sym(modifier_var)))
  
  # Require minimum sample size for reliable analysis
  if(nrow(metric_data) < 10) return(NULL)
  
  # Create sequence variable: Group 1 = cDBS→aDBS, Group 2 = aDBS→cDBS
  metric_data <- metric_data %>% mutate(Sequence = factor(Group))
  
  # BASE MODEL: Treatment effect without modifier interaction
  # Uses mixed-effects model with random intercept for each patient
  # to account for repeated measurements (same patient, different treatments)
  base_model <- tryCatch(
    lmer(Score_transformed ~ Treatment + Period + Sequence + Baseline_z + (1|PatientID),
         data = metric_data, REML = TRUE, control = lmerControl(optimizer = "bobyqa")),
    error = function(e) NULL
  )
  
  # MODIFIER MODEL: Adds interaction between Treatment and baseline modifier
  # The interaction term tells us if treatment effect varies by baseline value
  modifier_formula <- paste("Score_transformed ~ Treatment * ", modifier_var,
                            "+ Period + Sequence + Baseline_z + (1|PatientID)")
  
  modifier_model <- tryCatch(
    lmer(as.formula(modifier_formula), data = metric_data, REML = TRUE,
         control = lmerControl(optimizer = "bobyqa")),
    error = function(e) NULL
  )
  
  # Return NULL if either model fails to fit
  if(is.null(base_model) || is.null(modifier_model)) return(NULL)
  
  # Extract interaction effect coefficient (Treatment × Modifier)
  # This tells us how much the treatment effect changes per unit change in modifier
  model_summary <- summary(modifier_model)
  interaction_coef_name <- paste0("TreatmentaDBS:", modifier_var)
  
  if(!interaction_coef_name %in% rownames(model_summary$coefficients)) return(NULL)
  
  interaction_effect <- model_summary$coefficients[interaction_coef_name, ]
  
  # Calculate 95% confidence interval for interaction effect
  # Wald method provides approximate confidence intervals
  ci_result <- tryCatch({
    ci_manual <- confint(modifier_model, interaction_coef_name, method = "Wald")
    list(lower = ci_manual[1], upper = ci_manual[2])
  }, error = function(e) list(lower = NA, upper = NA))
  
  raw_estimate <- interaction_effect[1]
  
  # Apply unified sign convention: reverse sign for "lower is better" outcomes
  # This ensures positive = "worse baseline -> aDBS benefit" across all outcomes
  outcome_sign_multiplier <- if(metric %in% lower_the_better_metrics) -1 else 1
  
  unified_estimate <- raw_estimate * outcome_sign_multiplier
  unified_ci_lower_temp <- ci_result$lower * outcome_sign_multiplier
  unified_ci_upper_temp <- ci_result$upper * outcome_sign_multiplier
  # Ensure lower < upper after sign transformation
  unified_ci_lower <- min(unified_ci_lower_temp, unified_ci_upper_temp, na.rm = TRUE)
  unified_ci_upper <- max(unified_ci_lower_temp, unified_ci_upper_temp, na.rm = TRUE)
  
  # Calculate Cohen's d (standardized effect size)
  # This allows comparing effect sizes across different outcomes and modifiers
  # Effect size = interaction estimate / residual standard deviation
  residual_sd <- sigma(modifier_model)
  unified_effect_size <- unified_estimate / residual_sd
  
  # Calculate partial correlation (partial r)
  # This measures the strength of association between modifier and treatment effect,
  # while accounting for other variables in the model
  # Partial r ranges from -1 to +1, similar to regular correlation
  t_val <- interaction_effect[4]  # t-statistic
  df_val <- interaction_effect[3]  # degrees of freedom
  if(is.finite(t_val) && is.finite(df_val) && df_val > 0) {
    tcrit <- qt(0.975, df_val)  # Critical t-value for 95% CI
    # Convert t-statistic to partial correlation
    r_raw <- t_val / sqrt(t_val^2 + df_val)
    rL_raw <- (t_val - tcrit) / sqrt((t_val - tcrit)^2 + df_val)
    rU_raw <- (t_val + tcrit) / sqrt((t_val + tcrit)^2 + df_val)
    # Apply unified sign convention
    partial_r_unified <- outcome_sign_multiplier * r_raw
    partial_r_lower_tmp <- outcome_sign_multiplier * rL_raw
    partial_r_upper_tmp <- outcome_sign_multiplier * rU_raw
    partial_r_lower_unified <- min(partial_r_lower_tmp, partial_r_upper_tmp, na.rm = TRUE)
    partial_r_upper_unified <- max(partial_r_lower_tmp, partial_r_upper_tmp, na.rm = TRUE)
  } else {
    partial_r_unified <- partial_r_lower_unified <- partial_r_upper_unified <- NA
  }
  
  # Compare model fit: AIC (Akaike Information Criterion)
  # Lower AIC = better model. If modifier model has lower AIC, adding modifier
  # improves model fit (suggests modifier is important)
  aic_base <- AIC(base_model)
  aic_full <- AIC(modifier_model)
  
  # Clinical interpretation based on unified estimate
  # Positive: patients with worse baseline benefit more from aDBS
  # Negative: patients with better baseline benefit more from aDBS
  interpretation <- case_when(
    unified_estimate > 0 ~ "Worse baseline → aDBS superiority",
    unified_estimate < 0 ~ "Better baseline → aDBS superiority",
    TRUE ~ "No differential effect"
  )
  
  list(
    metric = metric,
    modifier = modifier_var,
    n_obs = nrow(metric_data),
    n_patients = length(unique(metric_data$PatientID)),
    interaction_estimate_raw = raw_estimate,
    ci_lower_raw = ci_result$lower,
    ci_upper_raw = ci_result$upper,
    interaction_estimate_unified = unified_estimate,
    interaction_se = interaction_effect[2],
    interaction_df = interaction_effect[3],
    interaction_t = interaction_effect[4],
    interaction_p_reference = interaction_effect[5],
    ci_lower_unified = unified_ci_lower,
    ci_upper_unified = unified_ci_upper,
    ci_width = unified_ci_upper - unified_ci_lower,
    effect_size_unified = unified_effect_size,
    ci_excludes_zero = (unified_ci_lower > 0 & unified_ci_upper > 0) | 
      (unified_ci_lower < 0 & unified_ci_upper < 0),
    partial_r_unified = partial_r_unified,
    partial_r_lower_unified = partial_r_lower_unified,
    partial_r_upper_unified = partial_r_upper_unified,
    aic_base = aic_base,
    aic_full = aic_full,
    aic_improvement = aic_base - aic_full,
    outcome_sign_multiplier = outcome_sign_multiplier,
    interpretation = interpretation,
    model_base = base_model,
    model_full = modifier_model
  )
}

# =============================================================================
# Run Analysis: Test All Modifier-Outcome Combinations
# =============================================================================
# Test each baseline predictor as a potential modifier for each outcome
# This creates a comprehensive exploration of which patient characteristics
# influence treatment response

# Collect available standardized modifier variables
modifier_vars <- c()
standardized_predictor_names <- paste0(predictor_vars, "_standardized")
for(pred_std in standardized_predictor_names) {
  if(pred_std %in% names(ancova_data)) {
    modifier_vars <- c(modifier_vars, pred_std)
  }
}
# Add sex if available (categorical variable, doesn't need standardization)
if("sex" %in% existing_predictors) {
  modifier_vars <- c(modifier_vars, "sex")
}

# Run analysis for all combinations
modifier_results <- list()
for(metric in eval_items_ordered) {
  for(modifier in modifier_vars) {
    result <- run_modifier_analysis(ancova_data, metric, modifier)
    if(!is.null(result)) {
      modifier_results[[paste(metric, modifier, sep = "_")]] <- result
    }
  }
}

# =============================================================================
# Compile Results: Create Summary Table
# =============================================================================
# Combine all analysis results into a single table for easy inspection
# Results are sorted by effect magnitude (largest absolute effects first)

modifier_summary <- map_dfr(modifier_results, function(result) {
  if(is.null(result)) return(NULL)
  
  tibble(
    Evaluation = result$metric,
    Modifier = result$modifier,
    N_Observations = result$n_obs,
    N_Patients = result$n_patients,
    Interaction_Estimate_Raw = result$interaction_estimate_raw,
    CI_Lower_Raw = result$ci_lower_raw,
    CI_Upper_Raw = result$ci_upper_raw,
    Interaction_Estimate_Unified = result$interaction_estimate_unified,
    Interaction_SE = result$interaction_se,
    Interaction_DF = result$interaction_df,
    Interaction_T = result$interaction_t,
    Interaction_P_reference = result$interaction_p_reference,
    CI_Lower_Unified = result$ci_lower_unified,
    CI_Upper_Unified = result$ci_upper_unified,
    CI_Width = result$ci_width,
    Effect_Size_Unified = result$effect_size_unified,
    Partial_r_Unified = result$partial_r_unified,
    Partial_r_Lower_Unified = result$partial_r_lower_unified,
    Partial_r_Upper_Unified = result$partial_r_upper_unified,
    CI_Excludes_Zero = result$ci_excludes_zero,
    AIC_Improvement = result$aic_improvement,
    Outcome_Sign_Multiplier = result$outcome_sign_multiplier,
    Interpretation = result$interpretation
  )
})

# =============================================================================
# Export Results and Summary
# =============================================================================

if(nrow(modifier_summary) > 0) {
  # Sort by absolute effect size (largest effects first)
  modifier_summary_ordered <- modifier_summary %>%
    arrange(desc(abs(Interaction_Estimate_Unified)))
  
  # Identify robust effects: confidence interval excludes zero
  # These are more reliable findings (less likely to be due to chance)
  robust_effects <- modifier_summary_ordered %>%
    filter(CI_Excludes_Zero == TRUE) %>%
    arrange(desc(Interaction_Estimate_Unified))
  
  write_csv(modifier_summary_ordered, "results_treatment_effect_modifiers.csv")
  saveRDS(modifier_results, "results_modifier_analysis_objects.rds")
  
  # Print Summary
  cat("\n=== ANALYSIS COMPLETE ===\n")
  cat(sprintf("Analyzed %d modifier-outcome combinations\n", nrow(modifier_summary)))
  cat(sprintf("Robust effects (CI excludes zero): %d\n", nrow(robust_effects)))
  
  if(nrow(robust_effects) > 0) {
    worse_baseline <- sum(robust_effects$Interaction_Estimate_Unified > 0)
    better_baseline <- sum(robust_effects$Interaction_Estimate_Unified < 0)
    cat(sprintf("- Worse baseline -> aDBS superiority: %d\n", worse_baseline))
    cat(sprintf("- Better baseline -> aDBS superiority: %d\n", better_baseline))
  }
  
  cat("\nExported files:\n")
  cat("- treatment_effect_modifiers.csv\n")
  cat("- modifier_analysis_objects.rds\n")
}